from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
import jieba

# 模拟邮件数据集
emails = [
    # 垃圾邮件样本 (标签为1)
    "免费获得 iPhone 点击链接 赢取大奖", 
    "赢取百万大奖立即参与 限时优惠", 
    "限时优惠不要错过 低价促销", 
    "恭喜中奖 领取奖金 点击链接", 
    "打折促销 买一送一 最后机会", 
    "低价商品 限量抢购 立即下单", 
    "免运费 特大优惠 不要错过", 
    "现金奖励 点击领取 优惠活动",
    # 正常邮件样本 (标签为0)  
    "明天会议安排请查收 项目讨论", 
    "项目报告已发送请审阅 工作总结", 
    "周末团建活动通知 团建安排", 
    "月度工作计划请查看 工作安排", 
    "财务报表已发送 请查收审核", 
    "下周日程安排 请确认时间", 
    "公司培训通知 请大家参加", 
    "项目进度汇报 请审阅修改"
]

labels = [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0] # 1:垃圾邮件, 0:正常邮件

# 自定义中文分词函数
def chinese_tokenizer(text):
    return jieba.lcut(text) # 精确模式分词

# 创建管道，指定分词器并尝试引入N-gram特征（如bigram）
model = make_pipeline(
    CountVectorizer(tokenizer=chinese_tokenizer, 
                    token_pattern=None,  # 禁用默认token_pattern
                    ngram_range=(1, 2)), # 同时考虑单个词和相邻的两个词组合
    MultinomialNB(alpha=0.5)  # 可尝试调整平滑强度，例如0.5
)
model.fit(emails, labels)

# 预测新邮件
test_emails = ["优惠活动进行中", "月度工作总结"]
predictions = model.predict(test_emails)
probabilities = model.predict_proba(test_emails)

print("完整概率矩阵:")
print(probabilities)
print("\n预测结果:")

for i, email in enumerate(test_emails):
    # 提取垃圾邮件的概率
    spam_prob = probabilities[i]
    
    print(f"\n邮件 '{email}' 的详细概率:")
    print(f"- 正常邮件概率: {spam_prob[0]:.2%}")
    print(f"- 垃圾邮件概率: {spam_prob[1]:.2%}")
    print(f"- 模型预测: {'垃圾邮件' if predictions[i] == 1 else '正常邮件'}")